import os


from typing import Dict
from utils.tools import read_config, logger, generate_chat_completion
from feapder.utils.tools import dumps_json
from app.lib.WeChatAlert import WeChatAlert
from app.client.MinioClient import minio_client
from app.client.LocalDocQaClient import LocalDocQAClient
from langchain.text_splitter import RecursiveCharacterTextSplitter
from funboost import (
    boost,
    BrokerEnum,
    BoosterParams,
    FunctionResultStatusPersistanceConfig,
)
from funboost_cli_user import project_root_path
from app.manager.BotManager import BotManager

"""
如果发布者要获取消费结果一定要设置is_using_rpc_mode=True
"""


@boost(
    BoosterParams(
        queue_name="wechat_alert_queue",
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        concurrent_num=1,
        is_using_rpc_mode=True,
        retry_interval=10,
        max_retry_times=10,
        is_show_message_get_from_broker=True,
        qps=1,
    )
)
def alert_wechat(wx_name="exlo", alert_message="这是一个测试告警"):
    wechat_alert = WeChatAlert()
    return wechat_alert.send_message(wx_name=wx_name, message=f"{alert_message}")


@boost(
    BoosterParams(
        queue_name="gen_a_queue",
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        concurrent_num=4,
        is_using_rpc_mode=True,
        retry_interval=10,
        max_retry_times=1,
        is_push_to_dlx_queue_when_retry_max_times=True,
        do_task_filtering=True,
        is_send_consumer_hearbeat_to_redis=True,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=2 * 365 * 7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
    )
)
def gen_a(question: str, knowledge_name: str) -> Dict:
    """
    封装的QA方法，用于查询文档并返回结果。

    :param question: 提问的问题
    :param knowledge_name: 知识库名称
    :return: 查询结果，以字典形式返回
    """

    # 读取配置文件
    try:
        c = read_config(knowledge_name)
    except Exception as e:
        logger.error(e)
        bot_manager =BotManager(knowledge_name)
        bot_info = bot_manager.init_bot_if_not_exists()
        bot_manager.save_bot_info(bot_info)
        c = read_config(knowledge_name)
        
    # 检查配置是否包含bot_id和kb_ids
    bot_id = c.get("bot_id")
    kb_ids = c.get("kb_ids")
    if not bot_id or not isinstance(kb_ids, list) or not kb_ids:
        raise Exception("配置信息不完整，缺少bot_id或kb_ids")

    # 初始化LocalDocQAClient
    dqac = LocalDocQAClient(question, kb_ids, bot_id=bot_id)

    # 执行查询
    res = dqac.query_local_doc()
    if not res:
        logger.error("查询结果为空")
        raise Exception("查询结果为空")

    # 记录查询结果
    logger.info(dumps_json(res))

    return res


@boost(
    BoosterParams(
        queue_name="gen_q_queue",
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        concurrent_num=10,
        is_send_consumer_hearbeat_to_redis=True,
        is_push_to_dlx_queue_when_retry_max_times=True,
        do_task_filtering=True,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
    )
)
def gen_q(txt, knowledge_name):
    """
    对切割的文本生成问题
        此处发送请求，使用chatgpt生成问题
        将生成的问题选择对应的知识库进行qa生成推送到gen_a
    """
    prompt = f"""
<参考文本来源于《{knowledge_name}》书>
<参考文本>
{txt}
</参考文本>
<指令>
- 请根据参考文本提出10个心理学相关的问题
- 不要出现第一人称
</指令>
    """
    response = generate_chat_completion(prompt)
    q_list = response.choices[0].message.content.split("\n")
    q_list = [_.split(".")[1].strip() for _ in q_list]
    for q in q_list:
        gen_a.push(q, knowledge_name)
    return response.to_dict()


@boost(
    BoosterParams(
        queue_name="split_txt_gen_q_queue",
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        is_send_consumer_hearbeat_to_redis=True,
        is_push_to_dlx_queue_when_retry_max_times=True,
        do_task_filtering=False,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
    )
)
def split_txt(file_name):
    knowledge_name = file_name.replace(".txt", "")
    txt = minio_client.get_txt("books", file_name)

    total = []
    for item in [(300, 100), (400, 150), (500, 200)]:
        chunk_size, chunk_overlap = item
        r_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,  # 块长度  # 重叠字符串长度
        )
        total.extend(r_splitter.split_text(txt))
    for txt in list(set(total)):
        gen_q.push(txt, knowledge_name)
    return len(total)


@boost(
    BoosterParams(
        queue_name="create_knowledge_queue",
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        is_send_consumer_hearbeat_to_redis=True,
        is_push_to_dlx_queue_when_retry_max_times=True,
        do_task_filtering=False,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
    )
)
def create_knowledge_base(filename):

    # 缓存到本地
    knowledge_name = filename.replace(".txt", "")
    path_dir = os.path.join(project_root_path, "data", knowledge_name)
    if not os.path.exists(path_dir):
        os.makedirs(path_dir)
        print("path_dir", path_dir)

    file_data = minio_client.download_data("books", filename)
    file_path = os.path.join(path_dir, filename)
    with open(file_path, "wb") as f:
        f.write(file_data)
    # 创建知识库
    knowledge_name = filename.replace(".txt", "")
    manager = BotManager(knowledge_name)
    flag = manager.run()
    # 删除文件
    os.remove(file_path)
    os.removedirs(path_dir)
    return flag

@boost(
    BoosterParams(
        queue_name='chatat_txt_to_qa_queue',
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        qps=10,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
        do_task_filtering=True,
    )
)
def chatat_txt_to_qa(book,paragraph):
    system_message = {
        "role": "system",
        "content": "把其中关键的内容提取出来，制作成我需要的对话数据集的格式，也就是常见的问答数据语料，请不要用第三视角制作数据集。请你用第一或者第二视角例如：问:你觉得为什么我总是倒霉？答:我觉得你太消极了，你可以多关注一些别的好的一方面。不要总关心坏的一方面。仿照我给的例子用问: 答: 这样的格式制作数据集。最后，问答对中，在答的部分，要给出足够长的回答。下面是原文："
    }
    prompt = paragraph

    response = generate_chat_completion(prompt=prompt,system_message=system_message)

    return  {'book':book,'qa_result': response.to_dict()}

@boost(
    BoosterParams(
        queue_name='chatat_process_text_queue',
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
        do_task_filtering=True,
    )
)

def process_text(book,sentences_per_paragraph=15):
 
    txt = minio_client.get_txt("books", book)
    content = txt.replace('\n', '')
    # 初始化变量
    paragraph = []
    sentence_count = 0
    paragraphs = []

    # 分割文本为句子并分组成段落
    for sentence in content.split('。'):
        if sentence:  # 忽略空句子
            paragraph.append(sentence + '。')
            sentence_count += 1
            if sentence_count >= sentences_per_paragraph:
                paragraphs.append('《问题》\n' + ''.join(paragraph))
                paragraph = []  # 重置段落
                sentence_count = 0

    # 添加最后一个段落（如果有）
    if paragraph:
        paragraphs.append('《问题》\n' + ''.join(paragraph))

    for paragraph in paragraphs:
        chatat_txt_to_qa.push(book,paragraph)
    return {'book':book,'paragraphs_count':len(paragraphs)}

@boost(
    BoosterParams(
        queue_name='chatat_chose_books_queue',
        broker_kind=BrokerEnum.RABBITMQ_AMQPSTORM,
        is_using_rpc_mode=True,
        function_result_status_persistance_conf=FunctionResultStatusPersistanceConfig(
            is_save_result=True,
            is_save_status=True,
            expire_seconds=7 * 24 * 3600,
            is_use_bulk_insert=True,
        ),
        do_task_filtering=True,
    )
)
def choice_book(book:str) -> str:
    """
    选择知识库
    :param book: 知识库名称
    :return: 知识库名称
    """
    
    process_text.push(book)